import os
import paramiko
import stat
import time
import psutil
import shutil


def scp(remote_host, remote_username, private_key, remote_dict, local_dict):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        key = paramiko.RSAKey.from_private_key_file(private_key)
        client.connect(remote_host, username=remote_username, pkey=key)
        scp_client = client.open_sftp()
        
        for day in os.listdir(local_dict):
            for clip in os.listdir(os.path.join(local_dict, day)):
                files = os.listdir(os.path.join(local_dict, day, clip))
                if not "updated" in files:
                    continue
                remote_data_dir = remote_dict + "/" + day  + "/" + clip
                # scp_client.mkdir(remote_data_dir)
                remote_vis_dir = remote_data_dir + "/vis"
                remote_color_dir = remote_vis_dir + "/color"
                remote_depth_dir = remote_vis_dir + "/depth"
                local_vis_dir = os.path.join(local_dict, day, clip, "out", "vis")
                if "vis" in files:
                    # only depth
                    scp_client.chdir(remote_depth_dir)
                    for old_depth in scp_client.listdir():
                        scp_client.remove(remote_depth_dir + "/" + old_depth)
                elif "depth" in files:
                    scp_client.mkdir(remote_vis_dir)
                    scp_client.mkdir(remote_color_dir)
                    scp_client.mkdir(remote_depth_dir)
                    for file in os.listdir(local_vis_dir):
                        if file.endswith(".json"):
                            scp_client.put(os.path.join(local_vis_dir, file), remote_vis_dir + "/" + file)
                    for color in os.listdir(os.path.join(local_vis_dir, "color")):
                        # print(os.path.join(local_vis_dir, "color", color))
                        scp_client.put(os.path.join(local_vis_dir, "color", color), remote_color_dir + "/" + color)
                for depth in os.listdir(os.path.join(local_vis_dir, "depth")):
                    scp_client.put(os.path.join(local_vis_dir, "depth", depth), remote_depth_dir + "/" + depth)
                scp_client.put(os.path.join(local_dict, day, clip, "updated"), remote_data_dir + "/" + "updated")
                shutil.rmtree(os.path.join(local_dict, day, clip))
                print(os.path.join(local_dict, day, clip))
                                
    except paramiko.AuthenticationException:
        print("Authentication Fail")
    except paramiko.SSHException as e:
        print("SSH:", str(e))
    finally:
        if scp_client:
            scp_client.close()
        if client:
            client.close()



if __name__ == "__main__":
    HOST = "10.210.5.10"
    USER = "liuyun"
    PK_KEY = "C:\\Users\\liuyu\\.ssh\\id_rsa"
    REMOTE_DIR = "/share/datasets/hhodataset/azure5"
    LOCAL_DIR = "D:\\mkv\\azure5"
    while(True):
        scp(HOST, USER, PK_KEY, REMOTE_DIR, LOCAL_DIR)
        time.sleep(5)
        